/*
  Python extension for the C interface to memcached, using apr_memcache library.
  Quick&Dirty hack from cmemcache.
  vcc (vcc.163@gmail.com)
*/

#include <Python.h>
#include "apr_memcache.h"

/*** Types ***/

typedef struct 
{
    PyObject_HEAD
    apr_pool_t *p;
    struct apr_memcache_t* mc;
    int debug;
} CmemcacheObject;

/*** Defines ***/

#ifdef NDEBUG
#define debug(args)
#define debug_def(args)
#else
#define debug(args) printf args
#define debug_def(args) args
#endif

#define MAX_SERVER 15

/*** Forward Declarations ***/

static void
cmemcache_dealloc(CmemcacheObject* self);

static char *
strnchr(const char *str, int c, const size_t len) {
  char *cp;
  size_t i;

  for (cp = str, i = 0; i < len && *cp != '\0'; i++, cp++) {
    if (c == (int)*cp)
      return cp;
  }

  return '\0';
}


//----------------------------------------------------------------------------------------
//
static int
do_set_servers(CmemcacheObject* self, PyObject* servers)
{
   int error = 0;
   const int size = PySequence_Size(servers);
   int i;
   debug(("do_set_servers\n"));
    
    if (!PySequence_Check(servers))
    {
        PyErr_BadArgument();
        return -1;
    }

    
    /* there seems to be no way to remove servers, so get rid of memcache all together */
    cmemcache_dealloc( self );
    assert(self->mc == NULL);

    /* create new instance */
    apr_pool_create(&self->p, NULL);
    apr_memcache_create(self->p, MAX_SERVER, 0, &self->mc);
    debug(("new mc %p\n", self->mc));
    if (self->mc == NULL)
    {
        PyErr_NoMemory();
        return -1;
    }
    
    /* add servers, allow any sequence of strings */
    
    for (i = 0; i < size && error == 0; ++i)
    {
        PyObject* item = PySequence_GetItem(servers, i);
        if (item)
        {
            PyObject* name = NULL;
            int weight = 1;
            
            if (PyString_Check(item))
            {
                name = item;
            }
            else if (PyTuple_Check(item))
            {
                error = ! PyArg_ParseTuple(item, "Oi", &name, &weight);
            }
            if (name)    
            {
                const char* cserver = PyString_AsString(name);
                assert(cserver);
                debug(("cserver %s weight %d\n", cserver, weight));
            
                /* mc_server_add4 is not happy without ':' (it segfaults!) so check */
                if (strstr(cserver, ":") == NULL)
                {
                    PyErr_Format(PyExc_TypeError,
                                 "expected \"server:port\" but \"%s\" found", cserver);
                    error = 1;
                }
                else
                {
                    int i;
                    if (weight>15)
                    {
                        weight = 15;
                    }
                    Py_BEGIN_ALLOW_THREADS;
                    for (i = 0; i < weight; ++i)
                    {
                    	apr_memcache_server_t* server = NULL;
                    	int len = strlen(cserver);
                    	char* cp = strnchr(cserver, ':', len);
                    	char* host = apr_pstrndup(self->p, cserver, cp-cserver);
                    	char* port = apr_pstrndup(self->p, cp+1, len-(cp-cserver));
                    	int port_i = atoi(port);
						int retval;
                    	if (port_i == 0) {
                    		PyErr_Format(PyExc_TypeError,
	                                 "port %s isn't integer", port);
	                    
	                    	error = 1;
	                    	break;
                    	}
                    	retval = apr_memcache_server_create(self->p, host, port_i, 0, 1, 1, 60, &server);
                    	retval = apr_memcache_add_server(self->mc, server);
                    	debug(("retval %d add server %s %d %p\n", retval, host, port_i, server));
                    }
                    Py_END_ALLOW_THREADS;
                }
            }
            else
            {
                PyErr_BadArgument();
                error = 1;
            }
            Py_DECREF(item);
        }
    }
    if (error)
    {
        apr_pool_destroy(self->p);
        self->p = NULL;
        self->mc = NULL;
        return -1;
    }
    return 0;
}

//----------------------------------------------------------------------------------------
//
static int
cmemcache_init(CmemcacheObject* self, PyObject* args, PyObject* kwds)
{
    PyObject* servers=NULL;
    char debug = 0;

    if (! PyArg_ParseTuple(args, "O|b", &servers, &debug))
        return -1; 

    self->debug = debug;
    self->mc = NULL;
    return do_set_servers(self, servers);
}

//----------------------------------------------------------------------------------------
//
static PyObject*
cmemcache_set_servers(PyObject* pyself, PyObject* servers)
{
    CmemcacheObject* self = (CmemcacheObject*)pyself;
    
    if (do_set_servers(self, servers) != -1)
    {
        Py_INCREF(Py_None);
        return Py_None;
    }
    /* fixme: what to return in the error case, shouldn't an exception have been raised? */
    return NULL;
}

//----------------------------------------------------------------------------------------
//
static void
cmemcache_dealloc(CmemcacheObject* self)
{
    debug(("cmemcache_dealloc\n"));
    
    if (self->mc)
    {
        Py_BEGIN_ALLOW_THREADS;
        if (self->p)
          apr_pool_destroy(self->p);
        self->p = NULL;
        self->mc = NULL;
        Py_END_ALLOW_THREADS;
    }
}

enum StoreType
{
    SET,
    ADD,
    REPLACE
};

//----------------------------------------------------------------------------------------
//
static PyObject*
cmemcache_store(PyObject* pyself, PyObject* args, enum StoreType storeType)
{
    char* key = NULL;
    int keylen = 0;
    const char* value = NULL;
    int valuelen = 0;
    int time = 0;
    int flags = 0;
    int retval = 0;
    CmemcacheObject* self = (CmemcacheObject*)pyself;
    
    assert(self->mc);
    
    if (! PyArg_ParseTuple(args, "s#s#|ii",
                           &key, &keylen, &value, &valuelen, &time, &flags))
        return NULL;
    
    
    Py_BEGIN_ALLOW_THREADS;
    debug(("cmemcache_store %d %s '%s' time %d flags %d\n",
           storeType, key, value, time, flags));
    switch(storeType)
    {
        case SET:
            retval = apr_memcache_set(self->mc, key, value, valuelen, time, flags);
            break;
        case ADD:
            retval = apr_memcache_add(self->mc, key, value, valuelen, time, flags);
            break;
        case REPLACE:
            retval = apr_memcache_replace(self->mc, key, value, valuelen, time, flags);
            break;
    }
    debug(("retval = %d\n", retval));
    Py_END_ALLOW_THREADS;
    
    return PyInt_FromLong(retval);
}

//----------------------------------------------------------------------------------------
//
static PyObject*
cmemcache_set(PyObject* pyself, PyObject* args)
{
    return cmemcache_store(pyself, args, SET);
}

//----------------------------------------------------------------------------------------
//
static PyObject*
cmemcache_add(PyObject* pyself, PyObject* args)
{
    return cmemcache_store(pyself, args, ADD);
}

//----------------------------------------------------------------------------------------
//
static PyObject*
cmemcache_replace(PyObject* pyself, PyObject* args)
{
    return cmemcache_store(pyself, args, REPLACE);
}

//----------------------------------------------------------------------------------------
//
static PyObject*
cmemcache_get_imp(PyObject* pyself, PyObject* args, int retFlags)
{
    char* key = NULL;
    int keylen = 0;
    PyObject* retval;
	int rt;
	apr_size_t len;
    char *result;
	apr_uint16_t res_flags; 
    CmemcacheObject* self = (CmemcacheObject*)pyself;
    
    assert(self->mc);

    if (! PyArg_ParseTuple(args, "s#", &key, &keylen))
    {
        debug(("bad arguments\n"));
        return NULL;
    }
    debug(("cmemcache_get_imp %s len %d\n", key, keylen));

	rt = apr_memcache_getp(self->mc, self->p, key, &result, &len, &res_flags);
	if (rt == APR_SUCCESS) {
		if (retFlags) {
			retval = Py_BuildValue("s#i", result, len, res_flags);
		} else {
			retval = PyString_FromStringAndSize(result, len);
		}
	}
    else
    {
        Py_INCREF(Py_None);
        retval = Py_None;
    }
    return retval;
}

//----------------------------------------------------------------------------------------
//
static PyObject*
cmemcache_get(PyObject* pyself, PyObject* args)
{
    return cmemcache_get_imp(pyself, args, 0);
}

//----------------------------------------------------------------------------------------
//
static PyObject*
cmemcache_getflags(PyObject* pyself, PyObject* args)
{
    return cmemcache_get_imp(pyself, args, 1);
}


//----------------------------------------------------------------------------------------
//
static PyObject*
cmemcache_delete(PyObject* pyself, PyObject* args)
{
    char* key = NULL;
    int keylen = 0;
    int time = 0;
    int retval;
    CmemcacheObject* self = (CmemcacheObject*)pyself;
    
    debug(("cmemcache_delete\n"));

    assert(self->mc);

    if (! PyArg_ParseTuple(args, "s#|i", &key, &keylen, &time))
        return NULL;

    
    Py_BEGIN_ALLOW_THREADS;
    debug(("cmemcache_delete %s time %d\n", key, time));
    retval = apr_memcache_delete(self->mc, key, time);
    debug(("retval = %d\n", retval));
    Py_END_ALLOW_THREADS;
    
    return PyInt_FromLong(retval);
}

//----------------------------------------------------------------------------------------
//
static PyObject*
cmemcache_incr_decr(PyObject* pyself, PyObject* args, int incr)
{
    char* key = NULL;
    int keylen = 0;
    int delta = 1;
    int newval;
	apr_status_t rt;

    CmemcacheObject* self = (CmemcacheObject*)pyself;
    
    debug(("cmemcache_incr_decr\n"));

    assert(self->mc);
    
    if (! PyArg_ParseTuple(args, "s#|i", &key, &keylen, &delta))
        return NULL;

    
    Py_BEGIN_ALLOW_THREADS;
	debug(("cmemcache_incr_decr %s %s delta %d\n", incr ? "incr" : "decr", key, delta));
    if ( incr )
    {
        rt = apr_memcache_incr(self->mc, key, delta, &newval);
    }
    else
    {
        rt = apr_memcache_decr(self->mc, key, delta, &newval);
    }
    Py_END_ALLOW_THREADS;

    if ( rt != APR_SUCCESS )
    {
        Py_INCREF(Py_None);
        return Py_None;
    }
    return PyInt_FromLong(newval);
}

//----------------------------------------------------------------------------------------
//
static PyObject*
cmemcache_incr(PyObject* pyself, PyObject* args)
{
    return cmemcache_incr_decr( pyself, args, 1 );
}

//----------------------------------------------------------------------------------------
//
static PyObject*
cmemcache_decr(PyObject* pyself, PyObject* args)
{
    return cmemcache_incr_decr( pyself, args, 0 );
}

//----------------------------------------------------------------------------------------
//
static PyObject*
cmemcache_get_stats(PyObject* pyself, PyObject* args)
{
	PyObject* retval = NULL;
	CmemcacheObject* self = (CmemcacheObject*)pyself;
	apr_memcache_stats_t* stats;
	int rt;
	int i;
	    
    debug(("cmemcache_get_stats %d \n", self->mc->ntotal));
    retval = PyList_New(0);

	for (i=0; i < self->mc->ntotal; i++) {
	    Py_BEGIN_ALLOW_THREADS;
		rt = apr_memcache_stats(self->mc->live_servers[i], self->p, &stats);
		Py_END_ALLOW_THREADS;
    	if (rt == APR_SUCCESS) {
            char buffer[128+1];
    		PyObject* name = NULL;
            PyObject* dict = PyDict_New();
            PyObject* tuple = PyTuple_New(2);
            snprintf(buffer, 128, "%s:%d", self->mc->live_servers[i]->host, self->mc->live_servers[i]->port);
    		name = PyString_FromString(buffer);
    
    #define SET_SNPRINTF(d, key, fmt, val)           \
            snprintf(buffer, 128, fmt, val);   \
            PyDict_SetItem(d, PyString_FromString(key), PyString_FromString(buffer))
            
    #define SET_TIME(d, key, val) SET_SNPRINTF(d, key, "%ld", val)
    #define SET_ITEMU32(d, key, val) SET_SNPRINTF(d, key, "%u", val)
    #define SET_ITEMU64(d, key, val) SET_SNPRINTF(d, key, "%llu", val)
    #define SET_TIMEVAL(d, key, val)                                \
            SET_SNPRINTF(d, key, "%lf", val)
                
            SET_ITEMU32(dict, "pid", stats->pid);
            SET_TIME(dict, "uptime", stats->uptime);
            SET_TIME(dict, "time", stats->time);
            PyDict_SetItem(dict, PyString_FromString("version"),
                           PyString_FromString(stats->version));
            SET_TIMEVAL(dict, "rusage_user", stats->rusage_user);
            SET_TIMEVAL(dict, "rusage_system", stats->rusage_system);
            SET_ITEMU32(dict, "curr_items", stats->curr_items);
            SET_ITEMU64(dict, "total_items", stats->total_items);
            SET_ITEMU64(dict, "bytes", stats->bytes);
            SET_ITEMU32(dict, "curr_connections", stats->curr_connections);
            SET_ITEMU64(dict, "total_connections", stats->total_connections);
            SET_ITEMU32(dict, "connection_structures", stats->connection_structures);
            SET_ITEMU64(dict, "cmd_get", stats->cmd_get);
            SET_ITEMU64(dict, "cmd_set", stats->cmd_set);
            SET_ITEMU64(dict, "get_hits", stats->get_hits);
            SET_ITEMU64(dict, "get_misses", stats->get_misses);
            SET_ITEMU64(dict, "bytes_read", stats->bytes_read);
            SET_ITEMU64(dict, "bytes_written", stats->bytes_written);
            SET_ITEMU64(dict, "limit_maxbytes", stats->limit_maxbytes);
    
            PyTuple_SetItem(tuple, 0, name);
            PyTuple_SetItem(tuple, 1, dict);
            PyList_Append(retval, tuple);
    	}
	}
	debug(("cmemcache_get_stats 2\n"));
	

	if (rt == APR_SUCCESS) {
        char buffer[128+1];
		PyObject* name = NULL;
        PyObject* dict = PyDict_New();
        PyObject* tuple = PyTuple_New(2);
        snprintf(buffer, 128, "%s:%s", "test", "11211");
		name = PyString_FromString(buffer);
        

#define SET_SNPRINTF(d, key, fmt, val)           \
        snprintf(buffer, 128, fmt, val);   \
        PyDict_SetItem(d, PyString_FromString(key), PyString_FromString(buffer))
        
#define SET_TIME(d, key, val) SET_SNPRINTF(d, key, "%ld", val)
#define SET_ITEMU32(d, key, val) SET_SNPRINTF(d, key, "%u", val)
#define SET_ITEMU64(d, key, val) SET_SNPRINTF(d, key, "%llu", val)
#define SET_TIMEVAL(d, key, val)                                \
        SET_SNPRINTF(d, key, "%lf", val)
            
        SET_ITEMU32(dict, "pid", stats->pid);
        SET_TIME(dict, "uptime", stats->uptime);
        SET_TIME(dict, "time", stats->time);
        PyDict_SetItem(dict, PyString_FromString("version"),
                       PyString_FromString(stats->version));
        SET_TIMEVAL(dict, "rusage_user", stats->rusage_user);
        SET_TIMEVAL(dict, "rusage_system", stats->rusage_system);
        SET_ITEMU32(dict, "curr_items", stats->curr_items);
        SET_ITEMU64(dict, "total_items", stats->total_items);
        SET_ITEMU64(dict, "bytes", stats->bytes);
        SET_ITEMU32(dict, "curr_connections", stats->curr_connections);
        SET_ITEMU64(dict, "total_connections", stats->total_connections);
        SET_ITEMU32(dict, "connection_structures", stats->connection_structures);
        SET_ITEMU64(dict, "cmd_get", stats->cmd_get);
        SET_ITEMU64(dict, "cmd_set", stats->cmd_set);
        SET_ITEMU64(dict, "get_hits", stats->get_hits);
        SET_ITEMU64(dict, "get_misses", stats->get_misses);
        SET_ITEMU64(dict, "bytes_read", stats->bytes_read);
        SET_ITEMU64(dict, "bytes_written", stats->bytes_written);
        SET_ITEMU64(dict, "limit_maxbytes", stats->limit_maxbytes);

        PyTuple_SetItem(tuple, 0, name);
        PyTuple_SetItem(tuple, 1, dict);
        PyList_Append(retval, tuple);
	}

  
    return retval;
}

static PyMethodDef cmemcache_methods[] = {
    {
        "set_servers", cmemcache_set_servers, METH_O,
        "set_servers(servers) -- set memcached servers"
    },
    
    {
        "set", cmemcache_set, METH_VARARGS,
        "set(key, value, time=0, flags=0) -- Unconditionally sets a key to a given value in the memcache.\n\n"
        "@return: Nonzero on success.\n@rtype: int\n"
    },
    
    {
        "add", cmemcache_add, METH_VARARGS,
        "add(key, value, time=0, flags=0) -- Add new key with value.\n\n"
        "Like L{set}, but only stores in memcache if the key doesn't already exist."
    },
    
    {
        "replace", cmemcache_replace, METH_VARARGS,
        "replace(key, value, time=0, flags=0) -- replace existing key with value.\n\n"
        "Like L{set}, but only stores in memcache if the key already exists.\n"
        "The opposite of L{add}."
    },
    
    {
        "get", cmemcache_get, METH_VARARGS,
        "get(key) -- Retrieves a key from the memcache.\n\n@return: The value or None."
    },
    
    {
        "getflags", cmemcache_getflags, METH_VARARGS,
        "getflags(key) -- Retrieves a key from the memcache.\n\n"
        "@return: The (value,flags) or None."
    },
    
    {
        "delete", cmemcache_delete, METH_VARARGS,
        "delete(key, time=0) -- Deletes a key from the memcache.\n\n"
        "@return: Nonzero on success.\n@rtype: int"
    },
    
    {
        "incr", cmemcache_incr, METH_VARARGS,
        "incr(key, delta=1)\n"
        "\n"
        "Sends a command to the server to atomically increment the value for C{key} by\n"
        "C{delta}, or by 1 if C{delta} is unspecified.  Returns None if C{key} doesn't\n"
        "exist on server, otherwise it returns the new value after incrementing.\n"
        "\n"
        "Note that the value for C{key} must already exist in the memcache, and it\n"
        "must be the string representation of an integer.\n"
        "\n"
        ">>> mc.set(\"counter\", \"20\")  # returns 1, indicating success\n"
        "1\n"
        ">>> mc.incr(\"counter\")\n"
        "21\n"
        ">>> mc.incr(\"counter\")\n"
        "22\n"
        "\n"
        "Overflow on server is not checked.  Be aware of values approaching\n"
        "2**32.  See L{decr}.\n"
        "\n"
        "@param delta: Integer amount to increment by (should be zero or greater).\n"
        "@return: New value after incrementing.\n"
        "@rtype: int or None if C{key} doesn't exist\n"
    },
    
    {
        "decr", cmemcache_decr, METH_VARARGS,
        "decr(key, delta=1)\n"
        "\n"
        "Like L{incr}, but decrements.  Unlike L{incr}, underflow is checked and\n"
        "new values are capped at 0.  If server value is 1, a decrement of 2\n"
        "returns 0, not -1.\n"
        "\n"
        "@param delta: Integer amount to decrement by (should be zero or greater).\n"
        "@return: New value after decrementing.\n"
        "@rtype: int or None if C{key} doesn't exist\n"
    },
    
    {
        "get_stats", cmemcache_get_stats, METH_NOARGS,
        "get_stats() -- Get statistics from all servers.\n"
        "@return: A list of tuples ( server_identifier, stats_dictionary ).\n"
        "The dictionary contains a number of name/value pairs specifying\n"
        "the name of the status field and the string value associated with\n"
        "it.  The values are not converted from strings."
    },
    
    {NULL}  /* Sentinel */
};

static PyTypeObject cmemcache_CmemcacheType = {
    PyObject_HEAD_INIT(NULL)
    0,                         /*ob_size*/
    "StringClient",            /*tp_name*/
    sizeof(CmemcacheObject),   /*tp_basicsize*/
    0,                         /*tp_itemsize*/
    (destructor)cmemcache_dealloc,         /*tp_dealloc*/
    0,                         /*tp_print*/
    0,                         /*tp_getattr*/
    0,                         /*tp_setattr*/
    0,                         /*tp_compare*/
    0,                         /*tp_repr*/
    0,                         /*tp_as_number*/
    0,                         /*tp_as_sequence*/
    0,                         /*tp_as_mapping*/
    0,                         /*tp_hash */
    0,                         /*tp_call*/
    0,                         /*tp_str*/
    0,                         /*tp_getattro*/
    0,                         /*tp_setattro*/
    0,                         /*tp_as_buffer*/
    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,        /*tp_flags*/
    "Client object",           /* tp_doc */
    0,		               /* tp_traverse */
    0,		               /* tp_clear */
    0,		               /* tp_richcompare */
    0,		               /* tp_weaklistoffset */
    0,		               /* tp_iter */
    0,		               /* tp_iternext */
    cmemcache_methods,         /* tp_methods */
    0,                         /* tp_members */
    0,                         /* tp_getset */
    0,                         /* tp_base */
    0,                         /* tp_dict */
    0,                         /* tp_descr_get */
    0,                         /* tp_descr_set */
    0,                         /* tp_dictoffset */
    (initproc)cmemcache_init,  /* tp_init */
    0,                         /* tp_alloc */
    PyType_GenericNew,         /* tp_new */
};

static PyMethodDef cmemcache_module_methods[] = {
    {NULL}  /* Sentinel */
};

#ifndef PyMODINIT_FUNC	/* declarations for DLL import/export */
#define PyMODINIT_FUNC void
#endif
PyMODINIT_FUNC
init_pyaprmemcache(void) 
{
    PyObject* m;

    debug(("init_pyaprmemcache\n"));
    
    cmemcache_CmemcacheType.tp_new = PyType_GenericNew;
    if (PyType_Ready(&cmemcache_CmemcacheType) < 0)
        return;

	apr_initialize();
    m = Py_InitModule3("_pyaprmemcache", cmemcache_module_methods,
                       "Extension to memcached using apr_memcache.");

    Py_INCREF(&cmemcache_CmemcacheType);
    PyModule_AddObject(m, "StringClient", (PyObject *)&cmemcache_CmemcacheType);
}

